import q0107
from astro.constants import wlya, Ckms
from astro.utilities import readatom, convolve_psf,calc_iontau
import astro.spec as sp

import pyfits
import numpy as np
import matplotlib.pyplot as pl
from matplotlib.pyplot import plot,hist,show

def plotmet(spec,ion,zp1,wavcand):
    """ Plot spectrum showing metal line candidates. """
    pl.ion()
    # show expected lines with sample N, b.

    #use smaller wave divisions so voigt profiles are sampled properly
    ndiv = 12  # brute force - should be a smarter way to do this!
    # there is a smarter way: use np.interp()
    swa = spec.wa
    # new array
    n = len(swa)
    wa = np.interp(np.linspace(0, n-1, (n-1)*ndiv), range(n), swa)
    
#     wa = np.empty(ndiv*len(swa) - (ndiv-1), np.float)
#     wa[::ndiv] = swa
#     for i in range(1,ndiv):
#         #print i,len(v1[i::ndiv])
#         wa[i::ndiv] = swa[:-1] + i*(swa[1:] - swa[:-1]) / float(ndiv)
    
    logN = 15.0
    b = 10.0
    tau = calc_iontau(wa,ion,zp1,logN,b)
    co0 = exp(-tau)

    # instrumental broadening
    if spec.vdelt < 0:
        raise Exception
    elif spec.vdelt < 45:
        # GHRS
        fwhm = 0.8 / (spec.vdelt / Ckms * speca.wa.mean())
    elif spec.vdelt >= 45:
        # FOS
        fwhm = 1.39 / (spec.vdelt / Ckms * speca.wa.mean())
    #print 'fwhm', fwhm
    co1 = convolve_psf(co0,fwhm*ndiv)

    # re-bin back to original wav array
    co2 = np.empty_like(swa)
    for i in range(len(swa)):
        co2[i] = co1[i*ndiv:i*ndiv+ndiv].mean() 

    co = spec.co * co2

    f = pl.figure(1)
    f.clf()
    #print o6a-10,o6b+10
    obswa = ion.wav*zp1
    wmin = obswa.min() - 10
    wmax = obswa.max() + 10
    spec.plot(limits=[wmin,wmax],show=0)
    pl.hold(1)
    i0,i1 = spec.wa_index(wmin), spec.wa_index(wmax)
    pl.gca().plot(spec.wa[i0:i1],co[i0:i1],'r',alpha=0.7)
    x0,x1,y0,y1 = pl.axis()
    pl.gca().vlines(obswa,0,y1*1.5,linestyle='dotted')
    # tick marks showing detected line positions
    i = np.array([spec.wa_index(w) for w in wavcand])
    ymin = spec.co[i]+(y1-y0)/10.
    ymax = spec.co[i]+(y1-y0)/4.
    pl.gca().vlines(wavcand,ymin,ymax,'k',linewidth=2)
    pl.gca().set_ylim(0,y1*1.5)
    pl.draw()


if 1:
    # minimum line detection significance.  Used for real and randoms.
    NSIGMA = 2.

    #########################################################################
    # Read in QSO, galaxy and absorption line data
    #########################################################################

    f = pyfits.open('ablines_fixed.fits')
    dat = f[1].data
    #rename the fields 'id' and 'field' to 'ID' and 'qfield' so the record
    #array doesn't get confused.
    dat = np.rec.fromrecords(dat.tolist(),names='lam_obs,lam_err,ew_obs,ew_err,'
                             'sigl,fwhm,ID,zobs,qfield,ew_rest,restlam,coldens,'
                             'rest_guess,zcalc,zdiff')

    # absorption lines.  This is the Morris linelist, which is taken from
    # Young et al. 01 and Dinshaw et al 97.

    mjla = dat[(dat.qfield == 1071)]         # Q0107a
    mjlb = dat[(dat.qfield == 1072)]         # Q0107b
    la_all = q0107.readabs('A')
    lb_all = q0107.readabs('B')
    lc_all = q0107.readabs('C')

    gal = q0107.readgal()                  # galaxies
    qso = q0107.readqsos()                 # qsos

    # want lines that are (1) Lya and (2) aren't too close to QSO Lya emission.

    maxwavs = wlya*(1+qso.z)*(1. - 1200./Ckms)   # 1200 km/s from lya emission
    lyalines = []
    for i,l in enumerate([la_all,lb_all,lc_all]):
        cond = np.array([item in ('unknown','lya') for item in l.ID])
        cond = cond & (l.wav < maxwavs[i])
        cond = cond & (l.signif_det > NSIGMA)
        temp = l[cond]
        temp.zabs = temp.wav/wlya - 1.
        lyalines.append(temp)

    la,lb,lc = lyalines


    # read in spectra
    speca0 = sp.read('fos/qsoa_hrs_uc.dat')
    specb0 = sp.read('fos/qsob_hrs_uc.dat')
    speca = sp.read('fos/qsoa_fos_uc.dat')
    specb = sp.read('fos/qsob_fos_uc.dat')
    specc = sp.read('fos/qsoc_uc.dat')

    # find potential c4, o6 lines, using lya positions as a basis.
    at = readatom('/home/nhmc/installed/vpfit/atom.dat')

    # make a fake transition, like CIV but with 3 Ang larger
    # wavelength separation.
    at['fake'] = at['OVI'].copy()
    i = at['fake'].wav.argmax()
    at['fake'].wav[i] = at['fake'].wav[i] + 8.0
    # tolerance in Angstroms to match lines
    tol = 2.0
    spec = dict(A0=speca0, B0=specb0, A=speca, B=specb, C=specc)

if 1:
    for lya,lall,los in zip([la,lb,lc],[la_all,lb_all,lc_all],'ABC'):
        zp1 = lya.wav / wlya
        ions = ['OVI','SiIV','CIV']
        #ions = ['AlIII']
        #ions = ['fake']
        for ion in ions:
            trans = at[ion]
            trans.sort(order='wav')
            trans_wa = np.expand_dims(trans.wav[:2],1) * zp1
            for j,(tra,trb) in enumerate(trans_wa.T):
                diffs = np.abs(tra - lall.wav)
                ia = diffs.argmin()
                #print inda, diffs[inda]
                if diffs[ia] > tol:  continue
                # otherwise check other doublet line
                diffs = np.abs(trb - lall.wav)
                ib = diffs.argmin()
                #print '**', indb, diffs[indb]
                if diffs[ib] > tol:  continue
                # otherwise we've found a candidate
                print 'LOS %s: %s candidate, z=%.5f,lya=%.2f' % (
                    los,ion,zp1[j]-1,wlya*zp1[j])
                print '  expected  wavs: %.2f %.2f dlam=%.2f' % (
                    tra,trb,trb-tra)
                print '  candidate wavs: %.2f %.2f dlam=%.2f' % (
                    lall.wav[ia],lall.wav[ib], lall.wav[ib]-lall.wav[ia])
                key = (los if trb > 1550 else los+'0')
                plotmet(spec[key],at[ion],zp1[j],[lall.wav[ia],lall.wav[ib]])
                raw_input('enter to continue')

